import os
import pandas as pd
import matplotlib.pyplot as plt

def plot_loss_curve(loss_file_path='resource/training_loss_improved.txt', output_path='resource/training_loss_improved.png'):
    """
    读取训练损失数据并绘制损失曲线图
    
    参数:
        loss_file_path: 损失数据文件路径
        output_path: 输出图像路径
    """
    # 检查文件是否存在
    if not os.path.exists(loss_file_path):
        print(f"错误: 找不到文件 {loss_file_path}")
        return
    
    # 读取损失数据
    try:
        # 使用pandas读取CSV文件
        df = pd.read_csv(loss_file_path)
        
        # 检查列名
        if 'Generator Loss' not in df.columns or 'Discriminator Loss' not in df.columns:
            print(f"错误: 文件格式不正确，需要包含'Generator Loss'和'Discriminator Loss'列")
            return
        
        # 绘制损失曲线
        plt.figure(figsize=(12, 6))
        plt.title("生成器和判别器训练损失曲线", fontsize=14)
        plt.plot(df['Generator Loss'], label="生成器(G)", color='blue', alpha=0.7)
        plt.plot(df['Discriminator Loss'], label="判别器(D)", color='red', alpha=0.7)
        plt.xlabel("迭代次数", fontsize=12)
        plt.ylabel("损失值", fontsize=12)
        plt.grid(True, linestyle='--', alpha=0.6)
        plt.legend(fontsize=12)
        
        # 保存图像
        plt.tight_layout()
        plt.savefig(output_path, dpi=300)
        print(f"损失曲线已保存到 {output_path}")
        
        # 显示图像
        plt.show()
        
    except Exception as e:
        print(f"处理文件时出错: {e}")

if __name__ == "__main__":
    plot_loss_curve()